"""
基于 CrewAI 的股票分析工具集
"""
from crewai_tools import BaseTool
from typing import Optional, Type, Any
from pydantic.v1 import BaseModel, Field
import sys
import os
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from service.news_kimi_das import get_sina_finance_news
from service.tushare_pool import getDayStock
from datetime import datetime, timedelta
import json
from colorama import Fore, Style


class StockDataInput(BaseModel):
    """股票数据工具的输入模型"""
    stock_code: str = Field(description="股票代码，如 600519.SH")
    days: int = Field(default=365, description="获取历史数据的天数")


class StockDataTool(BaseTool):
    """股票数据获取工具"""
    name: str = "股票数据工具"
    description: str = "获取指定股票的历史价格数据、成交量等信息"
    args_schema: Type[BaseModel] = StockDataInput
    
    def _run(self, stock_code: str, days: int = 365) -> str:
        """获取股票数据"""
        try:
            print(Fore.CYAN + f"    - [StockDataTool] 正在获取股票 {stock_code} 的数据...")
            
            today = datetime.now()
            start_date = (today - timedelta(days=days)).strftime("%Y%m%d")
            end_date = today.strftime("%Y%m%d")
            
            daily_k_data = getDayStock(
                stockCode=stock_code, 
                startDate=start_date, 
                endDate=end_date
            )
            
            if daily_k_data is None or daily_k_data.empty:
                return f"无法获取股票代码 {stock_code} 的数据"
            
            # 获取基本统计信息
            latest_data = daily_k_data.head(10)
            stats = {
                "最新收盘价": daily_k_data.iloc[0]["close"],
                "最高价": daily_k_data["high"].max(),
                "最低价": daily_k_data["low"].min(),
                "平均成交量": daily_k_data["vol"].mean(),
                "平均换手率": daily_k_data["turnover_rate"].mean() if "turnover_rate" in daily_k_data.columns else "N/A"
            }
            
            result = f"""股票 {stock_code} 数据报告：
            
统计信息：
- 最新收盘价: {stats['最新收盘价']}
- 区间最高价: {stats['最高价']}
- 区间最低价: {stats['最低价']}
- 平均成交量: {stats['平均成交量']:.0f}
- 平均换手率: {stats['平均换手率']}

最近10个交易日详细数据：
{latest_data.to_string()}
"""
            print(Fore.GREEN + f"    - [StockDataTool] 成功获取股票数据")
            return result
            
        except Exception as e:
            error_msg = f"获取股票数据时出错: {str(e)}"
            print(Fore.RED + f"    - [StockDataTool] {error_msg}")
            return error_msg


class NewsSearchInput(BaseModel):
    """新闻搜索工具的输入模型"""
    stock_code: str = Field(description="股票代码")
    query: Optional[str] = Field(default=None, description="自定义搜索查询")


class NewsSearchTool(BaseTool):
    """新闻搜索工具"""
    name: str = "新闻搜索工具"
    description: str = "搜索与指定股票相关的最新财经新闻和市场资讯"
    args_schema: Type[BaseModel] = NewsSearchInput
    
    def _run(self, stock_code: str, query: Optional[str] = None) -> str:
        """搜索新闻"""
        try:
            if query is None:
                # 根据股票代码生成查询
                stock_name = self._get_stock_name(stock_code)
                query = f"搜索关于{stock_name}（{stock_code}）的最新财经新闻"
            
            print(Fore.CYAN + f"    - [NewsSearchTool] 开始搜索股票 {stock_code} 的新闻...")
            print(Fore.CYAN + f"    - [NewsSearchTool] 查询内容: {query}")
            
            news_result = get_sina_finance_news(query=query)
            
            # 尝试解析JSON格式的新闻数据
            try:
                news_data = json.loads(news_result)
                if isinstance(news_data, dict) and "news_list" in news_data:
                    news_list = news_data["news_list"]
                    print(Fore.GREEN + f"    - [NewsSearchTool] 获取到 {len(news_list)} 条新闻")
                    
                    formatted_news = f"关于 {stock_code} 的最新新闻：\n"
                    formatted_news += "=" * 50 + "\n"
                    
                    for idx, news in enumerate(news_list[:5], 1):  # 限制为前5条
                        formatted_news += f"\n{idx}. {news.get('title', '无标题')}\n"
                        formatted_news += f"   时间: {news.get('time', '未知')}\n"
                        formatted_news += f"   内容摘要: {news.get('content', '无内容')[:200]}...\n"
                        formatted_news += "-" * 40 + "\n"
                    
                    return formatted_news
                else:
                    return f"关于 {stock_code} 的最新新闻：\n{news_result[:1000]}"
                    
            except json.JSONDecodeError:
                return f"关于 {stock_code} 的最新新闻：\n{news_result[:1000]}"
            
        except Exception as e:
            error_msg = f"搜索新闻时出错: {str(e)}"
            print(Fore.RED + f"    - [NewsSearchTool] {error_msg}")
            return error_msg
    
    def _get_stock_name(self, stock_code: str) -> str:
        """根据股票代码获取股票名称"""
        # 这里可以维护一个股票代码到名称的映射
        stock_names = {
            "600519.SH": "贵州茅台",
            "000001.SZ": "平安银行",
            "000002.SZ": "万科A",
            # 添加更多映射
        }
        return stock_names.get(stock_code, stock_code.split('.')[0])


class AnalysisInput(BaseModel):
    """分析工具的输入模型"""
    data: str = Field(description="需要分析的数据或文本")
    analysis_type: str = Field(default="comprehensive", description="分析类型：technical/fundamental/comprehensive")


class DataAnalysisTool(BaseTool):
    """数据分析工具"""
    name: str = "数据分析工具"
    description: str = "对股票数据进行技术分析和基本面分析"
    args_schema: Type[BaseModel] = AnalysisInput
    
    def _run(self, data: str, analysis_type: str = "comprehensive") -> str:
        """执行数据分析"""
        try:
            print(Fore.CYAN + f"    - [DataAnalysisTool] 开始进行{analysis_type}分析...")
            
            # 这里可以添加更复杂的分析逻辑
            # 比如计算技术指标、分析趋势等
            
            analysis_result = f"""
数据分析报告（{analysis_type}）：

原始数据摘要：
{data[:500]}...

分析结果：
1. 趋势分析：根据数据显示的价格走势，市场呈现一定的波动性
2. 成交量分析：成交量数据反映了市场的活跃程度
3. 技术指标：需要进一步计算MACD、RSI等指标
4. 风险评估：建议关注市场波动风险

分析时间：{datetime.now().strftime("%Y-%m-%d %H:%M:%S")}
"""
            print(Fore.GREEN + f"    - [DataAnalysisTool] 分析完成")
            return analysis_result
            
        except Exception as e:
            error_msg = f"数据分析时出错: {str(e)}"
            print(Fore.RED + f"    - [DataAnalysisTool] {error_msg}")
            return error_msg
